# -*- coding: utf-8 -*-
from __future__ import print_function
import sys
sys.path.append('/opt/alps/lib')
sys.path.append('/share/opt/alps/lib')
import pyalps
import matplotlib.pyplot as plt
import pyalps.plot
import numpy as np
import time
#import scipy.io
import glob, os, shutil
import multiprocessing

#----------------------1 Run the Task------------------------
a1 = time.time()
task_name = 'mps_BHmodel'  # global variable, can be accessed by function Runmps()
#sweeps = 30

def Runmps(para):
    (maxstates,sweeps) = para  ##infomation hidden
    parm = {}
    parm['LATTICE'] = 'open chain lattice'
    parm['L'] = 30
    parm['MODEL_LIBRARY'] = 'bhmodel.xml'
    parm['MODEL'] = 'my boson hubbard'
    parm['J'] = 1
    parm['U'] = 0.5
    parm['mu'] = -3
    parm['Nmax'] = 3
    parm['N_total'] = 30
    parm['CONSERVED_QUANTUMNUMBERS'] = 'N'
    parm['MAXSTATES'] = maxstates
    parm['SWEEPS'] = sweeps
    parm['NUMBER_EIGENVALUES'] = 1
    #    parm['MEASURE_LOCAL[local_density]'] = 'n'

    parm = [parm]
    input_file = pyalps.writeInputFiles(task_name + '_para=' + str(para), parm)
    pyalps.runApplication('mps_optim', input_file, writexml=True)

def delete_old_outfiles():  # 删除上一次运行得到的文件
    old_in_files = glob.glob('*%s*.in.*' % task_name)
    old_out_files = glob.glob('*%s*.out.*' % task_name)
    for i in old_in_files + old_out_files:
        try:
            os.remove(i)  #不能删文件夹
        except OSError:
            try:
                shutil.rmtree(i)  #用来删文件夹
            except:
                print('Something wrong when deleting files.')
delete_old_outfiles()  # 决定是否删除上一次运行得到的文件

pool = multiprocessing.Pool(processes=4)
Paras = [(150,10)]
print('report: This task has %d parallel task/tasks.' % (len(Paras)))
print('task %s is Running...' % (task_name))
pool.map(Runmps, Paras)
pool.close()
pool.join()
a2 = time.time()
print('run mps has spent %.1f seconds.' % (a2 - a1))

##如果不需要接着计算，删除占用大的MPS波函数chkp文件夹,就运行这个函数
def delete_MPS_wavefunction():
   chkp_files = glob.glob('*%s*.out.chkp' % task_name)
   for i in chkp_files:
       try:
           shutil.rmtree(i)  #用来删文件夹
       except:
           print('Maybe no such file to delete.')
   print('check point files deleted.')
##delete_MPS_wavefunction()  # 决定是否删除chkp文件夹

#----------------------2 结果分析----------------------------------------------
## 用spyder分析时请注释掉上面写runApplication的部分，这是我想到最好的办法了

result_files = pyalps.getResultFiles(prefix=task_name)
assert result_files != [], 'Report: Error, no result file has been calculated out. The calculation has definitly failed.'
## 如果没有产生结果文件，马上报错，而不是隐藏着继续运行
result_files.sort() #根据para值排好序,让图例排列合理
#result_files.insert(0,result_files[-1]) #对于这个特别情况，把50挪到最前面来
#del(result_files[-1])

# 适用多任务的收敛图,优化为适合不同sweeps了
iterations = pyalps.loadIterationMeasurements(
    [result_files[-1]], what=['Energy', 'TruncatedWeight']) #探知一个收敛数据的大小
Energy_vs_Iterations = pyalps.collectXY(iterations, x='iteration', y='Energy')
xdata = Energy_vs_Iterations[0].x  #这array里面的元素居然是string，难怪不好好排序的
(xsize,ysize)=(len(xdata),len(result_files))
xdatas=[]
ydatas=[]
for resfile in result_files:
    resfile = [resfile]
    iterations = pyalps.loadIterationMeasurements(
        resfile, what=['Energy', 'TruncatedWeight'])
    Energy_vs_Iterations = pyalps.collectXY(iterations, x='iteration', y='Energy')
    sweeps = int(Energy_vs_Iterations[0].props['SWEEPS'])
    xdata = Energy_vs_Iterations[0].x  #这array里面的元素居然是string，难怪不好好排序的
    xdata = np.array([int(i) for i in xdata])
    ydata = Energy_vs_Iterations[0].y  #这array里面的元素是标准的float64

    x_final = np.array([])
    y_final = np.array([])
    
    for nsweep in range(sweeps):
        mask = (xdata == nsweep)
        y_temp = ydata[mask]
        y_final = np.concatenate((y_final, y_temp))
        x_temp = np.linspace(nsweep, nsweep + 1, len(y_temp) + 1)
        x_temp = np.delete(x_temp, 0)
        x_final = np.concatenate((x_final, x_temp))
    xdatas.append(x_final)
    ydatas.append(y_final)
    
ydatas=[np.round(i,12) for i in ydatas]  ## 只保留MPS随机性下稳定的位数

U=float(0.5)  #选择加载U=几的标准数据
standard_data=np.load('BHmodelEg_Standard_U=%s.npz'%(U))
EgStandard=standard_data['Eg']
EgStandard=np.round(EgStandard,12)

label_list=[]  #从字符串中extract那个参数的数字,左分右分取中间
for i in result_files:
    i1=i.split('para=')[1]
    i2=i1.split('.task')[0]
    label_list.append(i2)

fig = plt.figure()
(fig, ax) = plt.subplots(figsize=(4, 3), dpi=300)
for i in range(len(result_files)): # 多个curve自动上色
    ax.plot(xdatas[i], ydatas[i]-EgStandard, label='D,sweeps=%s'%(label_list[i]))
ax.legend(loc=1,fontsize=8)
ax.set_xlabel('MPS sweeps')
ax.set_ylabel('Energy Difference with lux')
ax.set_yscale('log')
#ax.set_ylim([1e-8,1e2]) 
ax.grid()
ax.set_title('BHmodel MPS convergence')
plt.tight_layout() ##solve figure in pdf saved cutted off
fig.savefig('BHmdoel_sweeps_vs_energy.pdf')

#### 测量 n(i)
#eigen_measure_energy = pyalps.loadEigenstateMeasurements(result_files,'local_density')
##print(eigen_measure_energy)
#nn=eigen_measure_energy[0][0].y[0]
##print(nn)
#n_old=np.load('chongqi_nn.npz')
#print(nn==n_old['nn1'])
#np.savez('chongqi_nn.npz',nn1=nn)

### 能量
#print(result_files)
#eigen_measure_energy = pyalps.loadEigenstateMeasurements(result_files)
#print(eigen_measure_energy)
#Es = [i[0].y[0] for i in eigen_measure_energy]
#print(Es)
##print(Energies)
#print('These energies caled repeatedly by MPS are:\n%.12f\n%.12f\n%.12f\n%.12f\n' %
#      (Es[0], Es[1], Es[2], Es[3]))

## 保存为matlab的mat文件吧，这表示可以用熟悉的matlab处理
#scipy.io.savemat('MPSconvergency_history_noInteraction.mat', {
#    'MPS_sweeps': x_final,
#    'Energy_difference': y_final,
#    'low_standard':low_standard,
#    'model_name':'NoInteractionModel'
#})

### 画图先不谈了
#fig=plt.figure()
#(fig,ax)=plt.subplots(figsize=(4,3),dpi=100)
#ax.plot(x_final,y_final,'b')#,label='|Magneti|')
##ax.plot(x2,y2,'--',color='purple',label='Energy')
##ax.legend(loc=4)
#ax.set_xlabel('MPS sweeps')
#ax.set_ylabel('Energy')
#ax.set_yscale('log')
##ax.set_ylim([low_standard,1e-4])
#ax.grid()
#ax.set_title('No interaction delta=%s D=1500'%(delta))
##fig.savefig(task_name + '_sweeps_vs_energy.pdf')
#
